//      nmf_decentralized.cc
//        TODO: Decentralized NMF
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include "helper/stl_helper.h"
#include "helper/arma_helper.h"
#include "helper/mpi_helper.h"

#include "nmf.h"

extern po::variables_map vm;
extern mpi::communicator* p_world;

void build_graph(matrix_t& G, int M) {
  std::string graphFile = VM(vm, "graph", std::string);
  matrix_t bucky;
  read_sparse_matrix<matrix_t>(graphFile, bucky, false); // read graph from text file
  //TODO
}

void decentralized_nmf()
{
  RequiredParameter(vm, "ms")("input")("output")("log")("k")("alpha")("iter")("rho")("graph");
  reset_seed();

  // find rows assigned to each worker
  int M = p_world->size(); // number of workers
  std::vector<int> begin(M), end_exclusive(M); // rows assigned to each worker
  if (IS_MASTER_PROCESS) assign_rows(M, begin, end_exclusive);
  broadcast(*p_world, begin, 0);
  broadcast(*p_world, end_exclusive, 0);
  MSG("rank= "<< p_world->rank() << " begin=" << begin[p_world->rank()] 
      << " end=" << end_exclusive[p_world->rank()]);

  // build graph of worker nodes
  matrix_t G;
  vi_t neighbors;
  if (IS_MASTER_PROCESS) build_graph(G, M);
  broadcast(*p_world, G, 0);
  for (int j = 0; j < M; j++) if (G(PROCESS_RANK, j) > 0) neighbors.pb(j);
  
  // initialization at each worker
  matrix_t A_i;                // data
  matrix_t W_i, H_i;           // variables
  std::vector<matrix_t> Y_ij;  // dual variables
  std::vector<matrix_t> H_ij;  // store messages from neighbors
  read_data(A_i, begin[p_world->rank()], end_exclusive[p_world->rank()]);
  int m_i = A_i.n_rows, n = A_i.n_cols, k = VM(vm, "k", int);
  gen_sparse_matrix<matrix_t>(W_i, m_i, k, 0.5);
  gen_sparse_matrix<matrix_t>(H_i, k, n, 0.5);
  BOOST_FOREACH(int j, neighbors) {
    Y_ij.pb(arma::zeros<matrix_t>(k, n));
    H_ij.pb(arma::zeros<matrix_t>(k, n));
  }

  // read parameter from command line
  int maxiter = VM(vm, "iter", int);
  double alpha = VM(vm, "alpha", double);
  double rho = VM(vm, "rho", double);
  std::string inFile = VM(vm, "input", std::string);
  std::string outFile = VM(vm, "output", std::string);
  std::string logFile = VM(vm, "log", std::string);  

  // prepare temporaries
  matrix_t A_H(A_i.n_rows, H_i.n_rows);
  matrix_t A_W(A_i.n_cols, W_i.n_cols);
  matrix_t W_T_W(W_i.n_cols, W_i.n_cols);
  matrix_t H_T_H(H_i.n_rows, H_i.n_rows);

  matrix_t A_i_H(A_i.n_rows + k*sz(neighbors), A_i.n_cols);
  matrix_t W_i_H(W_i.n_rows + k*sz(neighbors), W_i.n_cols);
  A_i_H( arma::span(0, A_i.n_rows-1), arma::span::all ) = A_i;
  matrix_t Wnew = W_i, Hnew = H_i;
  

}
